#include <fcntl.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <unistd.h>
#include <stdint.h>
#include <iostream>
#include <atomic>
#include <emmintrin.h> // 用于_mm_pause指令

#include "cpupin.h"
#include "timestamp.h"

#define SHM_SIZE 64           // 小数据量，容易放入L1缓存
#define MSG_COUNT 1000000     // 消息传递次数

/**
 * 共享数据结构（设计为适合放入L1缓存）
 * 该结构体大小控制在64字节以内，以便完全驻留在CPU的L1缓存行中。
 * 使用volatile关键字防止编译器优化对flag的访问，确保内存可见性。
 */
typedef __attribute__((aligned(64))) struct {
  std::atomic<int> flag;                    // 同步标志：0表示数据未就绪，1表示数据就绪
  int seq_id;                               // 消息序列号，用于验证消息顺序
  uint64_t time_cycles;                     // 时间戳，记录写入时的CPU周期数
  char data[SHM_SIZE - sizeof(uint64_t) - sizeof(int) - sizeof(std::atomic<int>)]; // 实际数据负载
} SharedData;

/**
 * 线程参数结构体
 * 用于向线程传递共享内存指针和目标CPU ID。
 */
typedef struct {
  SharedData* shm;  // 指向共享内存区域的指针
  int cpu_id;       // 要绑定的CPU核心ID
} ThreadArgs;

// 线程函数：写入共享数据
void* writer_thread(void* arg) {
  ThreadArgs* args = (ThreadArgs*)arg;
  SharedData* shm = args->shm;
  // 将当前线程绑定到指定的CPU核心
  // 为了实现最低延迟的线程间通信，应将读写线程绑定到同一物理核心的两个逻辑核心（超线程）
  // 这样它们可以共享L1缓存，减少跨核同步开销
  cpupin(args->cpu_id);

  // 写入数据到共享内存（会被加载到L1缓存）
  for (int i = 0; i < MSG_COUNT; i++) {
    shm->seq_id = i;
    shm->time_cycles = rdtscp();  // 使用CPU周期计数代替系统时间
    std::atomic_thread_fence(std::memory_order_release); // 内存屏障确保数据写入完成
    shm->flag.store(1, std::memory_order_release);  // 通知读取方数据就绪
    while (shm->flag.load(std::memory_order_acquire) == 1) {
      _mm_pause(); // 减少CPU功耗，提示处理器其他线程可以使用超线程资源
    }
  }
  return NULL;
}

// 线程函数：读取共享数据
void* reader_thread(void* arg) {
  ThreadArgs* args = (ThreadArgs*)arg;
  SharedData* shm = args->shm;
  // 将当前线程绑定到指定的CPU核心
  // 为了实现最低延迟的线程间通信，应将读写线程绑定到同一物理核心的两个逻辑核心（超线程）
  // 这样它们可以共享L1缓存，减少跨核同步开销
  cpupin(args->cpu_id);
  uint64_t total_latency_cycles = 0;
  int mismatch_count = 0;

  // 从共享内存读取数据（直接从L1缓存获取）
  for (int i = 0; i < MSG_COUNT; i++) {
    while (shm->flag.load(std::memory_order_acquire) == 0) {
      _mm_pause(); // 减少CPU功耗，提示处理器其他线程可以使用超线程资源
    }
    std::atomic_thread_fence(std::memory_order_acquire); // 内存屏障确保数据读取正确
    uint64_t span_cycles = rdtscp() - shm->time_cycles;  // 使用CPU周期计数计算延迟
    total_latency_cycles += span_cycles;
    if (i != shm->seq_id) {
      mismatch_count++;
    }
    shm->flag.store(0, std::memory_order_release);  // 通知写入方已读取
  }

  double avg_cycles = static_cast<double>(total_latency_cycles) / MSG_COUNT;
  double avg_nanoseconds = cycles_to_nanoseconds(avg_cycles);
  double avg_microseconds = cycles_to_microseconds(avg_cycles);

  std::cout << "CPU Frequency: " << get_cpu_frequency() << " GHz" << std::endl;
  std::cout << "Average latency over " << MSG_COUNT << " messages: " 
            << avg_cycles << " CPU cycles" << std::endl;
  std::cout << "Average latency: " << avg_nanoseconds << " ns" << std::endl;
  std::cout << "Average latency: " << avg_microseconds << " μs" << std::endl;
  
  if (mismatch_count > 0) {
    std::cout << "Sequence mismatches: " << mismatch_count << std::endl;
  }
  return NULL;
}

int main(int argc, char** argv) {
  // 检查命令行参数数量
  if (argc < 3) {
    std::cerr << "Usage: " << argv[0] << " <reader_cpu_id> <writer_cpu_id>" << std::endl;
    std::cerr << "Example: " << argv[0] << " 8 9  # Binds threads to HT cores of the same physical core" << std::endl;
    return 1;
  }

  // 创建或打开一个命名共享内存对象
  // /ht_shm 是共享内存的名字，O_CREAT表示如果不存在则创建
  int fd = shm_open("/ht_shm", O_CREAT | O_RDWR, 0666);
  if (fd == -1) {
    perror("shm_open");
    return 1;
  }

  // 设置共享内存大小为SharedData结构体的大小
  ftruncate(fd, sizeof(SharedData));

  // 将共享内存映射到当前进程地址空间
  // MAP_SHARED保证其他进程/线程可见修改
  SharedData* shm = (SharedData*)mmap(NULL, sizeof(SharedData),
                                      PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
  if (shm == MAP_FAILED) {
    perror("mmap");
    close(fd);
    return 1;
  }

  shm->flag.store(0, std::memory_order_release);  // 初始化同步标志为未就绪状态

  // 创建线程参数，接收用户指定的CPU ID
  // 重要提示：为了获得最佳性能，请将reader和writer绑定到同一物理核心的两个逻辑核心（超线程对）
  // 例如，在支持超线程的系统上，核心i和i+1可能属于同一个物理核心
  ThreadArgs reader_args = {shm, atoi(argv[1])};
  ThreadArgs writer_args = {shm, atoi(argv[2])};

  // 创建读写线程
  pthread_t writer, reader;
  if (pthread_create(&reader, NULL, reader_thread, &reader_args) != 0) {
    std::cerr << "Failed to create reader thread" << std::endl;
    return 1;
  }
  if (pthread_create(&writer, NULL, writer_thread, &writer_args) != 0) {
    std::cerr << "Failed to create writer thread" << std::endl;
    return 1;
  }

  // 等待两个线程完成
  pthread_join(writer, NULL);
  pthread_join(reader, NULL);

  // 清理资源
  munmap(shm, sizeof(SharedData));  // 解除内存映射
  shm_unlink("/ht_shm");            // 删除共享内存对象

  return 0;
}